#pragma once

#include "ColorBlendAttachment.h"
#include "RHIPipeline.h"
#include "../Texture/ColorRenderTarget.h"
#include "../Texture/DepthRenderTarget.h"
#include "../Texture/ResolveRenderTarget.h"

class GraphicsPipeline : public RHIPipeline
{
	friend class RHIRenderPass;

	struct BuilderDetails {

		RefCountPtr<RHIDescriptor> descriptor;

		std::vector<RefCountPtr<ColorRenderTarget>> colorRenderTargets;
		std::vector<VkImageLayout> colorImageLayouts;
		RefCountPtr<DepthRenderTarget> depthRenderTarget;
		VkImageLayout depthImageLayout;
		RefCountPtr<ResolveRenderTarget> resolveRenderTarget;
		VkImageLayout resolveImageLayout;
		std::vector<RefCountPtr<RenderTarget>> inputRenderTargets;
		std::vector<VkImageLayout> inputImageLayouts;

		// The Option to wait the last pipeline's src Stage/Access finished
		VkPipelineStageFlags dstStageMask;
		VkAccessFlags dstAccessMask;

		// The option the next pipeline should wait until option finished
		VkPipelineStageFlags srcStageMask;
		VkAccessFlags srcAccessMask;

		// Every color renderTarget should have a colorBlend
		// size 0 -> use default, size 1 -> use the same one, size = rt size -> each one for each rt
		VkPipelineColorBlendStateCreateInfo colorBlendState;
		std::vector<RefCountPtr<ColorBlendAttachment>> colorBlendAttachments;

		VkPipelineTessellationStateCreateInfo tessellationState;

		// If no data, the viewportState still need to set count, and the dynamicStates should have data
		VkPipelineViewportStateCreateInfo viewportStateCreateInfo;
		std::vector<VkViewport> viewports;
		std::vector<VkRect2D> scissors;

		std::vector<VkDynamicState> dynamicStates;

		VkPipelineRasterizationStateCreateInfo rasterizationState;

		VkPipelineMultisampleStateCreateInfo multisampleState;

		VkPipelineDepthStencilStateCreateInfo depthStencilState;

		RefCountPtr<RHIPipeline> basePipelineHandle;
		int32 basePipelineIndex;

		VkRenderPass renderPass;

		BuilderDetails();
	};

public:

	class Builder : public BuilderBase<BuilderDetails> {
		friend struct BuilderDetails;
	public:
		Builder() noexcept;
		Builder(Builder const& rhs) noexcept;
		Builder(Builder&& rhs) noexcept;
		~Builder() noexcept;
		Builder& operator=(Builder const& rhs) noexcept;
		Builder& operator=(Builder&& rhs) noexcept;

		/************************************************************************/
		/*                              Descriptor                              */
		/************************************************************************/

		Builder& SetDescriptor(RefCountPtr<RHIDescriptor> descriptor) noexcept;

		/************************************************************************/
		/*                              RenderTarget                            */
		/************************************************************************/

		Builder& AddColorRenderTarget(RefCountPtr<ColorRenderTarget> renderTarget, VkImageLayout imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL) noexcept;
		Builder& SetDepthRenderTarget(RefCountPtr<DepthRenderTarget> renderTarget, VkImageLayout imageLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL) noexcept;
		Builder& SetResolveRenderTarget(RefCountPtr<ResolveRenderTarget> renderTarget, VkImageLayout imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL) noexcept;
		// All type of rendertarget could be input attachment
		Builder& AddInputRenderTarget(RefCountPtr<RenderTarget> renderTarget, VkImageLayout imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) noexcept;

		/************************************************************************/
		/*                         SubpassDependency                            */
		/************************************************************************/

		Builder& SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept;
		Builder& SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept;
		Builder& SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept;
		Builder& SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept;

		/************************************************************************/
		/*                           TessellationState                          */
		/************************************************************************/

		Builder& SetPatchControlPoints(uint32 patchControlPoints) noexcept;

		/************************************************************************/
		/*                           ViewportState                              */
		/************************************************************************/

		Builder& AddViewport(VkViewport viewport) noexcept;
		Builder& AddScissor(VkRect2D scissor) noexcept;

		/************************************************************************/
		/*                           DynamicState                               */
		/************************************************************************/

		Builder& AddDynamicState(VkDynamicState dynamicState, uint32 count = UINT32_MAX) noexcept;

		/************************************************************************/
		/*                           RasterizationState                         */
		/************************************************************************/

		Builder& SetDepthClampEnable(VkBool32 depthClampEnable) noexcept;
		Builder& SetRasterizerDiscardEnable(VkBool32 rasterizerDiscardEnable) noexcept;
		Builder& SetPolygonMode(VkPolygonMode polygonMode) noexcept;
		Builder& SetCullMode(VkCullModeFlags cullMode) noexcept;
		Builder& SetFrontFace(VkFrontFace frontFace) noexcept;
		Builder& SetDepthBiasEnable(VkBool32 depthBiasEnable) noexcept;
		Builder& SetDepthBiasConstantFactor(float depthBiasConstantFactor) noexcept;
		Builder& SetDepthBiasClamp(float depthBiasClamp) noexcept;
		Builder& SetDepthBiasSlopeFactor(float depthBiasSlopeFactor) noexcept;
		Builder& SetLineWidth(float lineWidth) noexcept;

		/************************************************************************/
		/*                           MultisampleState                           */
		/************************************************************************/

		Builder& SetRasterizationSamples(VkSampleCountFlagBits rasterizationSamples) noexcept;
		Builder& SetSampleShadingEnable(VkBool32 sampleShadingEnable) noexcept;
		Builder& SetMinSampleShading(float minSampleShading) noexcept;
		Builder& SetSampleMask(VkSampleMask* sampleMask) noexcept;
		Builder& SetAlphaToCoverageEnable(VkBool32 alphaToCoverageEnable) noexcept;
		Builder& SetAlphaToOneEnable(VkBool32 alphaToOneEnable) noexcept;

		/************************************************************************/
		/*                           DepthStencilState                          */
		/************************************************************************/

		Builder& SetDepthTestEnable(VkBool32 depthTestEnable) noexcept;
		Builder& SetDepthWriteEnable(VkBool32 depthWriteEnable) noexcept;
		Builder& SetDepthCompareOp(VkCompareOp depthCompareOp) noexcept;
		Builder& SetDepthBoundsTestEnable(VkBool32 depthBoundsTestEnable) noexcept;
		Builder& SetStencilTestEnable(VkBool32 stencilTestEnable) noexcept;
		Builder& SetMinDepthBounds(float minDepthBounds) noexcept;
		Builder& SetMaxDepthBounds(float maxDepthBounds) noexcept;

		Builder& SetFrontFailOp(VkStencilOp failOp) noexcept;
		Builder& SetFrontPassOp(VkStencilOp passOp) noexcept;
		Builder& SetFrontDepthFailOp(VkStencilOp depthFailOp) noexcept;
		Builder& SetFrontCompareOp(VkCompareOp compareOp) noexcept;
		Builder& SetFrontCompareMask(uint32 compareMask) noexcept;
		Builder& SetFrontWriteMask(uint32 writeMask) noexcept;
		Builder& SetFrontReference(uint32 reference) noexcept;

		Builder& SetBackFailOp(VkStencilOp failOp) noexcept;
		Builder& SetBackPassOp(VkStencilOp passOp) noexcept;
		Builder& SetBackDepthFailOp(VkStencilOp depthFailOp) noexcept;
		Builder& SetBackCompareOp(VkCompareOp compareOp) noexcept;
		Builder& SetBackCompareMask(uint32 compareMask) noexcept;
		Builder& SetBackWriteMask(uint32 writeMask) noexcept;
		Builder& SetBackReference(uint32 reference) noexcept;

		/************************************************************************/
		/*                           ColorBlendState                            */
		/************************************************************************/

		Builder& AddColorBlendAttachment(RefCountPtr<ColorBlendAttachment> colorBlendAttachment) noexcept;
		Builder& SetLogicOpEnable(VkBool32 logicOpEnable) noexcept;
		Builder& SetLogicOp(VkLogicOp logicOp) noexcept;
		Builder& SetBlendConstants(float b1, float b2, float b3, float b4) noexcept;

		/************************************************************************/
		/*                           BasePipeline                               */
		/************************************************************************/

		Builder& SetBasePipeline(RefCountPtr<RHIPipeline> basePipeline, int32 index = -1) noexcept;

		/************************************************************************/
		/*                           RenderPass                                 */
		/************************************************************************/

		Builder& SetRenderPass(VkRenderPass renderPass) noexcept;

		RefCountPtr<GraphicsPipeline> Build();
		RefCountPtr<GraphicsPipeline> Build(RenderProcess* process, const std::string& name);

	protected:

		void RemoveDynamicState(VkDynamicState dynamicState);

	private:

		friend class GraphicsPipeline;
	};

	GraphicsPipeline(BuilderDetails& builderDetails);

	virtual void ReleaseResource() override;

	void SetRenderPass(VkRenderPass renderPass, uint32 subpass);

	void Create();

protected:

	std::vector<RefCountPtr<ColorRenderTarget>> mColorRenderTargets;
	std::vector<VkImageLayout> mColorImageLayouts;
	RefCountPtr<DepthRenderTarget> mDepthRenderTarget;
	VkImageLayout mDepthImageLayout;
	RefCountPtr<ResolveRenderTarget> mResolveRenderTarget;
	VkImageLayout mResolveImageLayout;
	std::vector<RefCountPtr<RenderTarget>> mInputRenderTargets;
	std::vector<VkImageLayout> mInputImageLayouts;

	// The Option to wait the last pipeline's src Stage/Access finished
	VkPipelineStageFlags mDstStageMask;
	VkAccessFlags mDstAccessMask;

	// The option the next pipeline should wait until option finished
	VkPipelineStageFlags mSrcStageMask;
	VkAccessFlags mSrcAccessMask;

	VkPipelineColorBlendStateCreateInfo mColorBlendState;
	// Every color renderTarget should have a colorBlend
	// size 0 -> use default, size 1 -> use the same one, size = rt size -> each one for each rt
	std::vector<RefCountPtr<ColorBlendAttachment>> mColorBlendAttachments;

	VkPipelineTessellationStateCreateInfo mTessellationState;

	VkPipelineViewportStateCreateInfo mViewportStateCreateInfo;
	std::vector<VkViewport> mViewports;
	std::vector<VkRect2D> mScissors;

	std::vector<VkDynamicState> mDynamicStates;

	VkPipelineRasterizationStateCreateInfo mRasterizationState;

	VkPipelineMultisampleStateCreateInfo mMultisampleState;

	VkPipelineDepthStencilStateCreateInfo mDepthStencilState;

	RefCountPtr<RHIPipeline> mBasePipelineHandle;
	int32 mBasePipelineIndex;

	VkRenderPass mRenderPass;

	uint32 mSubpass;
};


#if 0

VkGraphicsPipelineCreateInfo pipelineCreateInfo{};
pipelineCreateInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;

std::vector<VkPipelineShaderStageCreateInfo> shadeStages;
for (const auto& shaderStage : mDescriptor->GetShaderStages())
shadeStages.push_back(shaderStage->GetShaderStageCreateInfo());
pipelineCreateInfo.stageCount = static_cast<uint32>(shadeStages.size());
pipelineCreateInfo.pStages = shadeStages.data();

pipelineCreateInfo.pVertexInputState = &mDescriptor->GetVertexInputStateCreateInfo();

pipelineCreateInfo.pInputAssemblyState = &mDescriptor->GetInputAssemblyStateCreateInfo();

pipelineCreateInfo.pTessellationState = builderDetails.tessellationState.sType == VK_STRUCTURE_TYPE_MAX_ENUM ? nullptr : &builderDetails.tessellationState;

VkPipelineViewportStateCreateInfo viewportStateCreateInfo{};
viewportStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
viewportStateCreateInfo.viewportCount = static_cast<uint32>(builderDetails.viewports.size());
viewportStateCreateInfo.pViewports = builderDetails.viewports.data();
viewportStateCreateInfo.scissorCount = static_cast<uint32>(builderDetails.scissors.size());
viewportStateCreateInfo.pScissors = builderDetails.scissors.data();
pipelineCreateInfo.pViewportState = &viewportStateCreateInfo;

pipelineCreateInfo.pRasterizationState = &builderDetails.rasterizationState;

pipelineCreateInfo.pMultisampleState = &builderDetails.multisampleState;

pipelineCreateInfo.pDepthStencilState = &builderDetails.depthStencilState;

std::vector<VkPipelineColorBlendAttachmentState> colorBlendAttachments;
for (auto& colorBlendAttachment : builderDetails.colorBlendAttachments)
colorBlendAttachments.push_back(colorBlendAttachment->GetColorBlendAttachment());
builderDetails.colorBlendState.attachmentCount = static_cast<uint32>(colorBlendAttachments.size());
builderDetails.colorBlendState.pAttachments = colorBlendAttachments.data();

pipelineCreateInfo.pColorBlendState = &builderDetails.colorBlendState;

VkPipelineDynamicStateCreateInfo dynamicStateCreateInfo{};
if (builderDetails.dynamicStates.size() > 0)
{
	dynamicStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
	dynamicStateCreateInfo.dynamicStateCount = static_cast<uint32>(builderDetails.dynamicStates.size());
	dynamicStateCreateInfo.pDynamicStates = builderDetails.dynamicStates.data();
	pipelineCreateInfo.pDynamicState = &dynamicStateCreateInfo;
}
else
{
	pipelineCreateInfo.pDynamicState = nullptr;
}

pipelineCreateInfo.layout = mDescriptor->GetPipelineLayout();

pipelineCreateInfo.renderPass = builderDetails.renderPass->GetRenderPass();

pipelineCreateInfo.subpass = builderDetails.subpass;

if (builderDetails.basePipelineHandle.IsValid())
{
	// 在基础管道之后创建的所有管道都将是衍生品
	pipelineCreateInfo.flags = VK_PIPELINE_CREATE_DERIVATIVE_BIT;
	// 基础管道将是我们创建的第一个管道
	pipelineCreateInfo.basePipelineHandle = builderDetails.basePipelineHandle->GetPipeline();
	// 只允许对基本管道使用句柄或索引
	// 当我们使用句柄时，我们必须将索引设置为 -1（参见规范的第 9.5 节）
	pipelineCreateInfo.basePipelineIndex = builderDetails.basePipelineIndex;
}
else
{
	pipelineCreateInfo.basePipelineHandle = VK_NULL_HANDLE;
	pipelineCreateInfo.basePipelineIndex = -1;
}

VK_CHECK_RESULT(vkCreateGraphicsPipelines(GVulkanRHI->GetLogicalDevice(), VK_NULL_HANDLE, 1, &pipelineCreateInfo, nullptr, &mPipeline))

#endif